import torch
from lightning import LightningModule

class DynamicQuantizerLitModule(LightningModule):

    def __init__(
        self,
        net: torch.nn.Module,
        compile: bool,
    ) -> None:
        """Initialize a `MNISTLitModule`.

        :param net: The model to train.
        :param optimizer: The optimizer to use for training.
        :param scheduler: The learning rate scheduler to use for training.
        """
        super().__init__()
        # this line allows to access init params with 'self.hparams' attribute
        # also ensures init params will be stored in ckpt

        self.net = net
